# Simulation of SEIR disease transmission model (for SARS epidemic)
# By Wentao Yang, Thomas Macht
# Based on:
#   Introduction to Computational Science -- Shiflet & Shiflet
#   Module 6.2
#   R solution by Stephen Davies, University of Mary Washington

import inspect
import os
import math
import random
import itertools
from matplotlib import gridspec
import numpy as np
import matplotlib.pyplot as plt
import datetime


def sars_no_actions():
    simulate_sars(cql_city_quarantine_limit_deaths_2weeks=10000000, qit_prob_quarantine_incoming_travellers=0.,
                  q_daily_quarantine_rate=0.)


def sars_quarantine_within_city_10():
    simulate_sars(cql_city_quarantine_limit_deaths_2weeks=10000000, qit_prob_quarantine_incoming_travellers=0.,
                  q_daily_quarantine_rate=0.1)


def sars_quarantine_within_city_20():
    simulate_sars(cql_city_quarantine_limit_deaths_2weeks=10000000, qit_prob_quarantine_incoming_travellers=0.,
                  q_daily_quarantine_rate=0.2)


def sars_quarantine_within_city_30():
    simulate_sars(cql_city_quarantine_limit_deaths_2weeks=10000000, qit_prob_quarantine_incoming_travellers=0.,
                  q_daily_quarantine_rate=0.3)


def sars_quarantine_within_city_40():
    simulate_sars(cql_city_quarantine_limit_deaths_2weeks=10000000, qit_prob_quarantine_incoming_travellers=0.,
                  q_daily_quarantine_rate=0.4)


def sars_quarantine_within_city():
    simulate_sars(cql_city_quarantine_limit_deaths_2weeks=10000000, qit_prob_quarantine_incoming_travellers=0.)


def sars_quarantine_cities_total():
    simulate_sars(qit_prob_quarantine_incoming_travellers=0.,
                  at_prob_may_travel_no_symptoms=0., ats_prob_may_travel_symptoms=0.)


def sars_quarantine_cities_partly_50_5():
    simulate_sars(qit_prob_quarantine_incoming_travellers=0.,
                  at_prob_may_travel_no_symptoms=0.5, ats_prob_may_travel_symptoms=0.05)


def sars_quarantine_cities_partly_70_10():
    simulate_sars(qit_prob_quarantine_incoming_travellers=0.,
                  at_prob_may_travel_no_symptoms=0.7, ats_prob_may_travel_symptoms=0.1)


def sars_quarantine_cities_partly_90_15():
    simulate_sars(qit_prob_quarantine_incoming_travellers=0.,
                  at_prob_may_travel_no_symptoms=0.9, ats_prob_may_travel_symptoms=0.15)


def sars_quarantine_cities_partly():
    simulate_sars(qit_prob_quarantine_incoming_travellers=0.)


def sars_incoming_quarantine_50():
    simulate_sars(qit_prob_quarantine_incoming_travellers=.5)


def sars_incoming_quarantine():
    simulate_sars(qit_prob_quarantine_incoming_travellers=1.)


def sars_no_actions_no_random():
    simulate_sars(cql_city_quarantine_limit_deaths_2weeks=10000000, qit_prob_quarantine_incoming_travellers=0.,
                  q_daily_quarantine_rate=0., use_random=False, num_runs=1)


def sars_quarantine_within_city_no_random():
    simulate_sars(cql_city_quarantine_limit_deaths_2weeks=10000000, qit_prob_quarantine_incoming_travellers=0.,
                  use_random=False, num_runs=1)


def sars_incoming_quarantine_no_random():
    simulate_sars(qit_prob_quarantine_incoming_travellers=1., use_random=False, num_runs=1)


def sars_quarantine_cities_partly_no_random():
    simulate_sars(qit_prob_quarantine_incoming_travellers=0., use_random=False, num_runs=1)


def sars():
    """This is the main entry point of the program.
    :return: None
    """
    sars_quarantine_cities_partly()


def simulate_sars(num_runs=50, use_random=True, **args):
    """
    Use this function to perform an experiment.
    This functions runs a simulation num_runs times.
    The detailed output (tables, graphs) is saved to SOURCE_DIR/results/results_TIMESTAMP (for later reference).
    In addition, a short summary is printed to the console and the output of the last run is also saved directly
    to the source directory (for quick access).
    :param num_runs: Number of runs
    :param args: Refer to simulate_sars_once for the remaining parameters
    :return: None
    """

    result_path = os.path.join(os.path.dirname(__file__), "results",
                               "results_" + datetime.datetime.now().strftime('%Y%m%d_%H%M%S'))

    # Save the values of D, quarantined and CQ of all runs.
    all_D = [None for _ in range(num_runs)]
    all_quarantined = [None for _ in range(num_runs)]
    all_CQ = [None for _ in range(num_runs)]

    # For each run ...
    for r in range(num_runs):
        print("Starting run " + str(r))

        # Create output directory for this run.
        run_path = os.path.join(result_path, "run" + str(r))
        os.makedirs(run_path)

        # Calculate the result for this run.
        result = simulate_sars_once(use_random=use_random, **args)

        # Number of cities
        num_cities = len(result["N"])
        # Initial population
        init_population = result["N"][0][0]
        
        col_formats = [".2f", "s", ".2f", ".2f", ".2f", ".2f", ".2f", ".2f", ".2f", ".2f", ".2f", ".2f", ".2f", ".2f", ".2f", ".2f", ".2f",
                           ".2f", ".2f", ".2f", ".2f", ".2f", ".2f", ".2f", ".2f", ".2f", ".2f"]
        
        if use_random:
            # We do not need to show decimals.
            col_formats = [".2f", "s", "d", "d", "d", "d", "d", "d", "d", "d", "d", "d", "d", "d", "d", "d", "d",
                           "d", "d", "d", "d", "d", "d", "d", "d", "d", "d"]

        # For each city ...
        for c in range(num_cities):
            paths = [run_path]

            # Write table to output files.
            write_city_table_file(paths, "city" + str(c) + ".dat",
                                  ["Day", "CQ?", "Incoming", "Outgoing", "Infected", "Quarantined",
                                   "N", "S", "SQ", "E", "EQ", "IU", "IQ", "ID", "D",
                                   "R", "S_in", "E_in", "IU_in", "R_in", "SQ_in", "EQ_in", "IQ_in", "S_out", "E_out",
                                   "IU_out", "R_out"],
                                  col_formats,
                                  [result['x'], result['CQ'][c], result['incoming'][c], result['outgoing'][c],
                                   result['infected'][c], result['quarantined'][c], result['N'][c],
                                   result['S'][c], result['SQ'][c], result['E'][c], result['EQ'][c], result['IU'][c],
                                   result['IQ'][c], result['ID'][c],
                                   result['D'][c], result['R'][c], result['S_in'][c], result['E_in'][c],
                                   result['IU_in'][c],
                                   result['R_in'][c], result['SQ_in'][c], result['EQ_in'][c], result['IQ_in'][c],
                                   result['S_out'][c], result['E_out'][c], result['IU_out'][c], result['R_out'][c]])

        # Create plot and save it to the output directory.
        plt = make_plot(["Population", "Susceptible", "Infected", "Quarantined", "Dead", "Recovered"],
                        [result['N'], result['S'], result['infected'], result['quarantined'], result['D'], result['R']],
                        init_population, result['x'])
        plt.savefig(os.path.join(result_path, "run" + str(r) + ".png"))

        # Save the results for D and quarantined.
        all_D[r] = result['D']
        all_quarantined[r] = result['quarantined']
        all_CQ[r] = result['CQ']

    # Print a summary to the console and save it to the output directory.
    make_summary(result_path, result["param_vals"], init_population, result['delta_x'], all_D, all_quarantined, all_CQ)

def simulate_sars_once(c_num_cities=5,
                       N0_init_population=1000000,
                       d_num_days=500,
                       i_iterations_per_day=1,
                       b_prob_transmission=.05,
                       k_num_daily_contacts=10,
                       m_daily_death_rate=.1,
                       p_daily_infection_rate=.2,
                       q_daily_quarantine_rate=.5,
                       u_daily_quarantine_leave_rate=.1,
                       v_daily_recovery_rate=.0625,
                       w_daily_detection_rate=.0375,
                       w2_daily_isolation_rate=.0375,
                       t_prob_wants_to_travel=.01,
                       at_prob_may_travel_no_symptoms=.3,
                       ats_prob_may_travel_symptoms=0.,
                       cql_city_quarantine_limit_deaths_2weeks=10,
                       tqc_prob_travel_to_quarantined_city=.1,
                       qit_prob_quarantine_incoming_travellers=0.,
                       use_random=True):
    """This function performs a single simulation run. To get reliable results, the output of this function
    should not be used directly. Instead, use simulate_sars to run the simulation multiple times.
    :param c_num_cities: Number of cities
    :type c_num_cities: int
    :param N0_init_population: Initial population per city
    :type N0_init_population: int
    :param d_num_days: Number of days
    :type d_num_days: int
    :param i_iterations_per_day: Number of iterations per day
    :type i_iterations_per_day: int
    :param b_prob_transmission: Probability of transmission
    :type b_prob_transmission: float
    :param k_num_daily_contacts: Number of contacts per day
    :type k_num_daily_contacts: int
    :param m_daily_death_rate: Deaths per infected person per day
    :type m_daily_death_rate: float
    :param p_daily_infection_rate: New infections per exposed person per day
    :type p_daily_infection_rate: float
    :param q_daily_quarantine_rate: Exposed individuals quarantined per day
    :type q_daily_quarantine_rate: float
    :param u_daily_quarantine_leave_rate: Quarantined individuals leaving quarantine per day
    :type u_daily_quarantine_leave_rate: float
    :param v_daily_recovery_rate: Recoveries per infection per day
    :type v_daily_recovery_rate: float
    :param w_daily_detection_rate: Undetected infections detected per day
    :type w_daily_detection_rate: float
    :param w2_daily_isolation_rate: Quarantined infectees isolated per day
    :type w2_daily_isolation_rate: float
    :param t_prob_wants_to_travel: Probability that a person wants to travel
    :type t_prob_wants_to_travel: float
    :param at_prob_may_travel_no_symptoms: Probability that a person without symptoms is allowed travel
    :type at_prob_may_travel_no_symptoms: float
    :param ats_prob_may_travel_symptoms: Probability that a person with symptoms is allowed to travel
    :type ats_prob_may_travel_symptoms: float
    :param cql_city_quarantine_limit_deaths_2weeks: A city is quarantined when at least this number of individuals died
     within the last two weeks. This means that transportation from this city is restricted.
    :type cql_city_quarantine_limit_deaths_2weeks: int
    :param tqc_prob_travel_to_quarantined_city: Probability that individuals who want to travel to a quarantined city.
    :type tqc_prob_travel_to_quarantined_city: float
    :param qit_prob_quarantine_incoming_travellers: Probability that incoming travellers are quarantined (if the source
     city is quarantined)
    :type qit_prob_quarantine_incoming_travellers: float
    :param use_random: If true, do not use fractions of individuals but transform them into full individuals
     with a random generator.
    :type use_random: bool
    :return: None
    """

    params, _, _, vals = inspect.getargvalues(inspect.currentframe())
    param_vals = "\n".join([_p + " = " + str(vals[_p]) for _p in params])

    # Check types of parameters.
    assert type(c_num_cities) == int
    assert type(N0_init_population) == int
    assert type(d_num_days) == int
    assert type(i_iterations_per_day) == int
    assert type(b_prob_transmission) == float
    assert type(k_num_daily_contacts) == int
    assert type(m_daily_death_rate) == float
    assert type(p_daily_infection_rate) == float
    assert type(q_daily_quarantine_rate) == float
    assert type(u_daily_quarantine_leave_rate) == float
    assert type(v_daily_recovery_rate) == float
    assert type(w_daily_detection_rate) == float
    assert type(w2_daily_isolation_rate) == float
    assert type(t_prob_wants_to_travel) == float
    assert type(at_prob_may_travel_no_symptoms) == float
    assert type(ats_prob_may_travel_symptoms) == float
    assert type(cql_city_quarantine_limit_deaths_2weeks) == int
    assert type(tqc_prob_travel_to_quarantined_city) == float
    assert type(qit_prob_quarantine_incoming_travellers) == float
    assert type(use_random) == bool

    # Check validity of parameters.
    assert not m_daily_death_rate + v_daily_recovery_rate > 1  # Deaths + recoveries per infection per day
    assert not m_daily_death_rate + v_daily_recovery_rate + w2_daily_isolation_rate > 1  # Deaths + recoveries per infection + quarantined infectees isolated per day
    assert not m_daily_death_rate + v_daily_recovery_rate + w_daily_detection_rate > 1  # Deaths + recoveries per infection + undetected infections detected per day

    # Time increment
    delta_x = 1. / i_iterations_per_day

    # Vector of time values
    x = [_u * delta_x for _u in range(d_num_days * i_iterations_per_day + 1)]
    num_iterations = len(x)  # Number of iterations

    initInfected = 1.  # Number of initially infected and undetected individuals

    # Set up stock variables.
    S = init_two_dim_list(c_num_cities, num_iterations, 0)
    SQ = init_two_dim_list(c_num_cities, num_iterations, 0)
    E = init_two_dim_list(c_num_cities, num_iterations, 0)
    EQ = init_two_dim_list(c_num_cities, num_iterations, 0)
    IU = init_two_dim_list(c_num_cities, num_iterations, 0)
    IQ = init_two_dim_list(c_num_cities, num_iterations, 0)
    ID = init_two_dim_list(c_num_cities, num_iterations, 0)
    D = init_two_dim_list(c_num_cities, num_iterations, 0)
    R = init_two_dim_list(c_num_cities, num_iterations, 0)

    # Total population per city
    N = init_two_dim_list(c_num_cities, num_iterations, 0)

    # Saves whether transportation was shutdown per city and iteration
    CQ = init_two_dim_list(c_num_cities, num_iterations, False)
    # Probability to travel to a city (1 normally, different if transportation shutdown)
    # per city and iteration
    PTT = init_two_dim_list(c_num_cities, num_iterations, 1)

    # People who want to travel per city and iteration
    SWT = init_two_dim_list(c_num_cities, num_iterations, 0)
    EWT = init_two_dim_list(c_num_cities, num_iterations, 0)
    IUWT = init_two_dim_list(c_num_cities, num_iterations, 0)
    RWT = init_two_dim_list(c_num_cities, num_iterations, 0)

    # People who are allowed to travel per city and iteration
    SAT = init_two_dim_list(c_num_cities, num_iterations, 0)
    EAT = init_two_dim_list(c_num_cities, num_iterations, 0)
    IUAT = init_two_dim_list(c_num_cities, num_iterations, 0)
    RAT = init_two_dim_list(c_num_cities, num_iterations, 0)

    # Outgoing travellers per source, iteration and destination (fractions)
    S_out_perDest = [[[0 for _ in range(c_num_cities)] for _ in range(num_iterations)] for _ in range(c_num_cities)]
    E_out_perDest = [[[0 for _ in range(c_num_cities)] for _ in range(num_iterations)] for _ in range(c_num_cities)]
    IU_out_perDest = [[[0 for _ in range(c_num_cities)] for _ in range(num_iterations)] for _ in range(c_num_cities)]
    R_out_perDest = [[[0 for _ in range(c_num_cities)] for _ in range(num_iterations)] for _ in range(c_num_cities)]

    # Outgoing travellers per source, iteration and destination (integers)
    S_out_perDest_int = [[[0 for _ in range(c_num_cities)] for _ in range(num_iterations)] for _ in range(c_num_cities)]
    E_out_perDest_int = [[[0 for _ in range(c_num_cities)] for _ in range(num_iterations)] for _ in range(c_num_cities)]
    IU_out_perDest_int = [[[0 for _ in range(c_num_cities)] for _ in range(num_iterations)] for _ in range(c_num_cities)]
    R_out_perDest_int = [[[0 for _ in range(c_num_cities)] for _ in range(num_iterations)] for _ in range(c_num_cities)]

    # Incoming travellers per destination and iteration (at the beginning of the iteration)
    S_in = init_two_dim_list(c_num_cities, num_iterations, 0)
    E_in = init_two_dim_list(c_num_cities, num_iterations, 0)
    IU_in = init_two_dim_list(c_num_cities, num_iterations, 0)
    R_in = init_two_dim_list(c_num_cities, num_iterations, 0)
    SQ_in = init_two_dim_list(c_num_cities, num_iterations, 0)
    EQ_in = init_two_dim_list(c_num_cities, num_iterations, 0)
    IQ_in = init_two_dim_list(c_num_cities, num_iterations, 0)

    # Outgoing travellers per destination and iteration (at the end of the iteration)
    S_out = init_two_dim_list(c_num_cities, num_iterations, 0)
    E_out = init_two_dim_list(c_num_cities, num_iterations, 0)
    IU_out = init_two_dim_list(c_num_cities, num_iterations, 0)
    R_out = init_two_dim_list(c_num_cities, num_iterations, 0)

    # Initialize populations
    for i in range(0, c_num_cities):
        N[i][0] = N0_init_population

    # initInfected infected persons in city 0
    S[0][0] = N[0][0] - initInfected
    IU[0][0] = initInfected
    # No infected persons in all the other cities
    for i in range(1, c_num_cities):
        S[i][0] = N[i][0]

    # For each iteration do ...
    for i in range(1, len(x)):

        # For every city do ...
        for j in range(0, c_num_cities):
            contacts_with_IU = S[j][i - 1] *k_num_daily_contacts * IU[j][i - 1] / N[j][i - 1]

            SQ_to_S = SQ[j][i - 1] * u_daily_quarantine_leave_rate
            EQ_to_IQ = EQ[j][i - 1] * p_daily_infection_rate
            E_to_IU = E[j][i - 1] * p_daily_infection_rate
            ID_to_D = ID[j][i - 1] * m_daily_death_rate
            ID_to_R = ID[j][i - 1] * v_daily_recovery_rate
            IQ_to_D = IQ[j][i - 1] * m_daily_death_rate
            IQ_to_ID = IQ[j][i - 1] * w2_daily_isolation_rate
            IQ_to_R = IQ[j][i - 1] * v_daily_recovery_rate
            IU_to_D = IU[j][i - 1] * m_daily_death_rate
            IU_to_ID = IU[j][i - 1] * w_daily_detection_rate
            IU_to_R = IU[j][i - 1] * v_daily_recovery_rate
            S_to_E = contacts_with_IU * (1 - q_daily_quarantine_rate) * b_prob_transmission
            S_to_EQ = contacts_with_IU * q_daily_quarantine_rate * b_prob_transmission
            S_to_SQ = contacts_with_IU * q_daily_quarantine_rate * (1 - b_prob_transmission)

            if use_random:
                # Convert fractions of individuals to integer individuals.
                # None of the states must have a negative number of individuals.
                SQ_to_S_int = min(get_num_individuals_from_fractions(SQ_to_S * delta_x), SQ[j][i - 1])
                EQ_to_IQ_int = min(get_num_individuals_from_fractions(EQ_to_IQ * delta_x), EQ[j][i - 1])
                E_to_IU_int = min(get_num_individuals_from_fractions(E_to_IU * delta_x), E[j][i - 1])
                ID_to_D_int = min(get_num_individuals_from_fractions(ID_to_D * delta_x), ID[j][i - 1])
                ID_to_R_int = min(get_num_individuals_from_fractions(ID_to_R * delta_x), ID[j][i - 1] - ID_to_D_int)
                IQ_to_D_int = min(get_num_individuals_from_fractions(IQ_to_D * delta_x), IQ[j][i - 1])
                IQ_to_ID_int = min(get_num_individuals_from_fractions(IQ_to_ID * delta_x), IQ[j][i - 1] - IQ_to_D_int)
                IQ_to_R_int = min(get_num_individuals_from_fractions(IQ_to_R * delta_x),
                                  IQ[j][i - 1] - IQ_to_D_int - IQ_to_ID_int)
                IU_to_D_int = min(get_num_individuals_from_fractions(IU_to_D * delta_x), IU[j][i - 1])
                IU_to_ID_int = min(get_num_individuals_from_fractions(IU_to_ID * delta_x), IU[j][i - 1] - IU_to_D_int)
                IU_to_R_int = min(get_num_individuals_from_fractions(IU_to_R * delta_x),
                                  IU[j][i - 1] - IU_to_D_int - IU_to_ID_int)
                S_to_E_int = min(get_num_individuals_from_fractions(S_to_E * delta_x), S[j][i - 1])
                S_to_EQ_int = min(get_num_individuals_from_fractions(S_to_EQ * delta_x), S[j][i - 1] - S_to_E_int)
                S_to_SQ_int = min(get_num_individuals_from_fractions(S_to_SQ * delta_x),
                                  S[j][i - 1] - S_to_E_int - S_to_EQ_int)

                S_prime = SQ_to_S_int - S_to_SQ_int - S_to_E_int - S_to_EQ_int
                SQ_prime = S_to_SQ_int - SQ_to_S_int
                E_prime = S_to_E_int - E_to_IU_int
                EQ_prime = S_to_EQ_int - EQ_to_IQ_int
                IU_prime = E_to_IU_int - IU_to_R_int - IU_to_ID_int - IU_to_D_int
                IQ_prime = EQ_to_IQ_int - IQ_to_ID_int - IQ_to_D_int - IQ_to_R_int
                ID_prime = IU_to_ID_int + IQ_to_ID_int - ID_to_D_int - ID_to_R_int
                D_prime = IQ_to_D_int + IU_to_D_int + ID_to_D_int
                R_prime = IU_to_R_int + ID_to_R_int + IQ_to_R_int
            else:
                S_prime = SQ_to_S - S_to_SQ - S_to_E - S_to_EQ
                SQ_prime = S_to_SQ - SQ_to_S
                E_prime = S_to_E - E_to_IU
                EQ_prime = S_to_EQ - EQ_to_IQ
                IU_prime = E_to_IU - IU_to_R - IU_to_ID - IU_to_D
                IQ_prime = EQ_to_IQ - IQ_to_ID - IQ_to_D - IQ_to_R
                ID_prime = IU_to_ID + IQ_to_ID - ID_to_D - ID_to_R
                D_prime = IQ_to_D + IU_to_D + ID_to_D
                R_prime = IU_to_R + ID_to_R + IQ_to_R

                S_prime *= delta_x
                SQ_prime *= delta_x
                E_prime *= delta_x
                EQ_prime *= delta_x
                IU_prime *= delta_x
                IQ_prime *= delta_x
                ID_prime *= delta_x
                D_prime *= delta_x
                R_prime *= delta_x

            # Increase or decrease the population of each category based on
            # the rates of change for this time step.
            S[j][i] = S[j][i - 1] + S_prime
            SQ[j][i] = SQ[j][i - 1] + SQ_prime
            E[j][i] = E[j][i - 1] + E_prime
            EQ[j][i] = EQ[j][i - 1] + EQ_prime
            IU[j][i] = IU[j][i - 1] + IU_prime
            IQ[j][i] = IQ[j][i - 1] + IQ_prime
            ID[j][i] = ID[j][i - 1] + ID_prime
            D[j][i] = D[j][i - 1] + D_prime
            R[j][i] = R[j][i - 1] + R_prime

            # The population size must stay the same
            total_old = S[j][i - 1] + SQ[j][i - 1] + E[j][i - 1] + EQ[j][i - 1] + IU[j][i - 1] + IQ[j][i - 1] + ID[j][
                i - 1] + D[j][i - 1] + R[j][i - 1]
            total_new = S[j][i] + SQ[j][i] + E[j][i] + EQ[j][i] + IU[j][i] + IQ[j][i] + ID[j][i] + D[j][i] + R[j][i]

            if use_random:
                assert total_old == total_new

            # No category must be negative
            assert S[j][i] >= 0
            assert SQ[j][i] >= 0
            assert E[j][i] >= 0
            assert EQ[j][i] >= 0
            assert IU[j][i] >= 0
            assert IQ[j][i] >= 0
            assert ID[j][i] >= 0
            assert D[j][i] >= 0
            assert R[j][i] >= 0

            # Determine number of individuals who want to travel
            SWT[j][i] = S[j][i] * t_prob_wants_to_travel * delta_x
            EWT[j][i] = E[j][i] * t_prob_wants_to_travel * delta_x
            IUWT[j][i] = IU[j][i] * t_prob_wants_to_travel * delta_x
            RWT[j][i] = R[j][i] * t_prob_wants_to_travel * delta_x

            deaths = D[j][i] - D[j][max(0, i - 13)]

            # Transportation is limited
            if deaths >= cql_city_quarantine_limit_deaths_2weeks:
                CQ[j][i] = True
                PTT[j][i] = tqc_prob_travel_to_quarantined_city

                # Determine number of individuals who are allowed to travel
                SAT[j][i] = SWT[j][i] * at_prob_may_travel_no_symptoms
                EAT[j][i] = EWT[j][i] * at_prob_may_travel_no_symptoms
                # These individuals show symptoms, so the probability they are allowed to
                # travel is lower.
                IUAT[j][i] = IUWT[j][i] * ats_prob_may_travel_symptoms
                RAT[j][i] = RWT[j][i] * at_prob_may_travel_no_symptoms
            else:
                # Everyone who wants to travel is allowed to
                SAT[j][i] = SWT[j][i]
                EAT[j][i] = EWT[j][i]
                IUAT[j][i] = IUWT[j][i]
                RAT[j][i] = RWT[j][i]

            # For every city do ...
            for h in range(0, c_num_cities):
                # For every other city do ...
                if h != j:
                    # Divide the travellers over all other cities
                    S_out_perDest[j][i][h] = SAT[j][i] / (c_num_cities - 1) * PTT[h][i - 1]
                    E_out_perDest[j][i][h] = EAT[j][i] / (c_num_cities - 1) * PTT[h][i - 1]
                    IU_out_perDest[j][i][h] = IUAT[j][i] / (c_num_cities - 1) * PTT[h][i - 1]
                    R_out_perDest[j][i][h] = RAT[j][i] / (c_num_cities - 1) * PTT[h][i - 1]

                    if use_random:
                        # Fractions of individuals cannot travel
                        # Calculate the number of individuals who actually travel
                        # S, E, IU and R must not become negative
                        S_out_perDest_int[j][i][h] = min(get_num_individuals_from_fractions(S_out_perDest[j][i][h]),
                                                         S[j][i] - S_out[j][i])
                        E_out_perDest_int[j][i][h] = min(get_num_individuals_from_fractions(E_out_perDest[j][i][h]),
                                                         E[j][i] - E_out[j][i])
                        IU_out_perDest_int[j][i][h] = min(get_num_individuals_from_fractions(IU_out_perDest[j][i][h]),
                                                          IU[j][i] - IU_out[j][i])
                        R_out_perDest_int[j][i][h] = min(get_num_individuals_from_fractions(R_out_perDest[j][i][h]),
                                                         R[j][i] - R_out[j][i])

                        # Add individuals who are allowed to travel to the outgoing states
                        S_out[j][i] += S_out_perDest_int[j][i][h]
                        E_out[j][i] += E_out_perDest_int[j][i][h]
                        IU_out[j][i] += IU_out_perDest_int[j][i][h]
                        R_out[j][i] += R_out_perDest_int[j][i][h]
                    else:
                        # Add individuals who are allowed to travel to the outgoing states
                        S_out[j][i] += S_out_perDest[j][i][h]
                        E_out[j][i] += E_out_perDest[j][i][h]
                        IU_out[j][i] += IU_out_perDest[j][i][h]
                        R_out[j][i] += R_out_perDest[j][i][h]

            # Subtract outgoing individuals from population
            S[j][i] -= S_out[j][i]
            E[j][i] -= E_out[j][i]
            IU[j][i] -= IU_out[j][i]
            R[j][i] -= R_out[j][i]

        # Totals for checks
        total_in = 0
        total_out = 0
        total_N = 0

        # For every city do ...
        for j in range(0, c_num_cities):

            # For every source city do
            for h in range(0, c_num_cities):
                # Transportation in source city shutdown
                if CQ[h][i]:
                    # Add incoming travellers
                    if use_random:
                        S_in[j][i] += S_out_perDest_int[h][i][j] * (1 - qit_prob_quarantine_incoming_travellers)
                        E_in[j][i] += E_out_perDest_int[h][i][j] * (1 - qit_prob_quarantine_incoming_travellers)
                        IU_in[j][i] += IU_out_perDest_int[h][i][j] * (1 - qit_prob_quarantine_incoming_travellers)
                        R_in[j][i] += R_out_perDest_int[h][i][j]
                        # Add incoming travellers to quarantine
                        SQ_in[j][i] += S_out_perDest_int[h][i][j] * qit_prob_quarantine_incoming_travellers
                        EQ_in[j][i] += E_out_perDest_int[h][i][j] * qit_prob_quarantine_incoming_travellers
                        IQ_in[j][i] += IU_out_perDest_int[h][i][j] * qit_prob_quarantine_incoming_travellers
                    else:
                        S_in[j][i] += S_out_perDest[h][i][j] * (1 - qit_prob_quarantine_incoming_travellers)
                        E_in[j][i] += E_out_perDest[h][i][j] * (1 - qit_prob_quarantine_incoming_travellers)
                        IU_in[j][i] += IU_out_perDest[h][i][j] * (1 - qit_prob_quarantine_incoming_travellers)
                        R_in[j][i] += R_out_perDest[h][i][j]
                        # Add incoming travellers to quarantine
                        SQ_in[j][i] += S_out_perDest[h][i][j] * qit_prob_quarantine_incoming_travellers
                        EQ_in[j][i] += E_out_perDest[h][i][j] * qit_prob_quarantine_incoming_travellers
                        IQ_in[j][i] += IU_out_perDest[h][i][j] * qit_prob_quarantine_incoming_travellers
                else:
                    # Add incoming travellers
                    if use_random:
                        S_in[j][i] += S_out_perDest_int[h][i][j]
                        E_in[j][i] += E_out_perDest_int[h][i][j]
                        IU_in[j][i] += IU_out_perDest_int[h][i][j]
                        R_in[j][i] += R_out_perDest_int[h][i][j]
                    else:
                        S_in[j][i] += S_out_perDest[h][i][j]
                        E_in[j][i] += E_out_perDest[h][i][j]
                        IU_in[j][i] += IU_out_perDest[h][i][j]
                        R_in[j][i] += R_out_perDest[h][i][j]

            # Add incoming travellers
            S[j][i] += S_in[j][i]
            E[j][i] += E_in[j][i]
            IU[j][i] += IU_in[j][i]
            R[j][i] += R_in[j][i]
            SQ[j][i] += SQ_in[j][i]
            EQ[j][i] += EQ_in[j][i]
            IQ[j][i] += IQ_in[j][i]

            # Calculate new total population
            N[j][i] = S[j][i] + SQ[j][i] + E[j][i] + EQ[j][i] + IU[j][i] + \
                      IQ[j][i] + ID[j][i] + D[j][i] + R[j][i]

            # Totals for checks
            total_in += S_in[j][i] + E_in[j][i] + IU_in[j][i] + R_in[j][i] + SQ_in[j][i] + EQ_in[j][i] + IQ_in[j][i]
            total_out += S_out[j][i] + E_out[j][i] + IU_out[j][i] + R_out[j][i]
            total_N += N[j][i]

        if use_random:
            # The total population must remain the same
            assert total_N == N0_init_population * c_num_cities
            # The number of ingoing and outgoing travellers must be equal
            assert total_in == total_out

    # Calculate sums    
    incoming = (np.matrix(S_in) + np.matrix(E_in) + np.matrix(IU_in) + np.matrix(R_in) +
                np.matrix(SQ_in) + np.matrix(EQ_in) + np.matrix(IQ_in)).tolist()
    outgoing = (np.matrix(S_out) + np.matrix(E_out) + np.matrix(IU_out) +
                np.matrix(R_out)).tolist()
    infected = (np.matrix(IU) + np.matrix(IQ) + np.matrix(ID) + np.matrix(E) +
                np.matrix(EQ)).tolist()
    quarantined = (np.matrix(EQ) + np.matrix(IQ) + np.matrix(SQ) +
                   np.matrix(ID)).tolist()

    return dict(CQ=CQ, N=N, S=S, SQ=SQ, E=E, EQ=EQ, IU=IU, IQ=IQ, ID=ID, D=D, R=R, S_in=S_in, E_in=E_in, IU_in=IU_in,
                R_in=R_in, SQ_in=SQ_in, EQ_in=EQ_in, IQ_in=IQ_in, S_out=S_out, E_out=E_out, IU_out=IU_out, R_out=R_out,
                incoming=incoming, outgoing=outgoing, infected=infected, quarantined=quarantined, x=x,
                param_vals=param_vals, delta_x=delta_x)


def make_summary(result_path, param_vals, init_population, delta_x, D_per_run_per_city, quarantined_per_run_per_city, CQ_per_run_per_city):
    """Prints a summary of the simulation to the console and saves it result_path/summary.txt.
    :param result_path: Path for the result file
    :type result_path: String
    :param param_vals: Values of all the parameters
    :type param_vals: String
    :param init_population: The initial population
    :type init_population: int
    :param delta_x: Time increment
    :type delta_x: float (<= 1)
    :param D_per_run_per_city: D values per run per city
    :type D_per_run_per_city: List of List of List float
    :param quarantined_per_run_per_city: quarantined values per run per city
    :type quarantined_per_run_per_city: List of List of List of float
    :param CQ_per_run_per_city: CQ values per run per city
    :type CQ_per_run_per_city: List of List of List of bool
    :return: None
    """

    num_runs = len(D_per_run_per_city)
    num_cities = len(D_per_run_per_city[0])
    num_rows = len(D_per_run_per_city[0][0])

    total_init_population = init_population * num_cities
    days = num_rows * delta_x

    total_affected = 0
    total_deaths = 0
    total_lost = 0
    total_CQ = 0
    total_quarantined = 0

    total_summary = "\n### Parameters\n" + param_vals + "\n"

    # For each run ...
    for r in range(num_runs):
        affected = 0
        deaths = 0

        # For each city ...
        for D in D_per_run_per_city[r]:
            # Get the final number of deaths.
            deaths += D[num_rows - 1]
            # A city is considered affected if at least 5 % of the population died.
            if D[num_rows - 1] >= init_population * 0.05:
                affected += 1

        lost = np.sum(np.sum(np.matrix(quarantined_per_run_per_city[r]) + np.matrix(D_per_run_per_city[r])))
        CQ_city = list(itertools.chain(*CQ_per_run_per_city[r]))
        CQ = sum([1 for _cq in CQ_city if _cq])
        quarantined = np.sum(np.sum(np.matrix(quarantined_per_run_per_city[r])))

        total_affected += affected
        total_deaths += deaths
        total_lost += lost
        total_CQ += CQ
        total_quarantined += quarantined

        total_summary += prepare_summary_string("### Run " + str(r), affected, num_cities, deaths, total_init_population, lost, days, CQ, quarantined)

    avg_affected = total_affected / float(num_runs)
    avg_deaths = total_deaths / float(num_runs)
    avg_lost = total_lost / float(num_runs)
    avg_CQ = total_CQ / float(num_runs)
    avg_quarantined = total_quarantined / float(num_runs)

    total_summary += prepare_summary_string("### Total", avg_affected, num_cities, avg_deaths, total_init_population, avg_lost, days, avg_CQ, avg_quarantined)

    # Print summary to console.
    print total_summary

    # Write summary to file.
    outfile = open(os.path.join(result_path, "summary.txt"), 'w+')
    outfile.write(total_summary)
    outfile.close()


def prepare_summary_string(header, affected, num_cities, deaths, total_init_population, lost, days, total_CQ, total_quarantined):
    """Prepares a summary string.
    :param header: Text to be put in the beginning
    :type header: String
    :param affected: Affected cities (i.e. cities where at least 5% of the population died)
    :type affected: float
    :param num_cities: Number of cities
    :type num_cities: int
    :param deaths: Total number of deaths
    :type deaths: float
    :param total_init_population: Total initial population
    :type total_init_population: int
    :param lost: Total lost person days (i.e. days where a person was dead, quarantined or isolated)
    :type lost: float
    :param days: Total person days (i.e. days * initial population)
    :type days    :param total_CQ: Total days cities where quarantined
    :type days: float
    :param total_CQ: Total city quarantine days
    :type total_CQ: float
    :param total_quarantined: Total individuals in quarantine
    :type total_quarantined: float
    :return: Summary string
    """

    summary = "\n" + header
    summary += "\n" + "Cities with at least 5% of N0 dead:".ljust(40, ' ') + "%.2f (%.2f %%)" % (
        affected, affected / float(num_cities) * 100)
    summary += "\n" + "Total deaths:".ljust(40, ' ') + "%.2f (%.2f %%)" % (deaths, deaths / float(total_init_population) * 100)
    summary += "\n" + "Total lost days:".ljust(40, ' ') + "%.2f (%.2f %%)" % (lost, lost / float(days) / float(num_cities) / float(total_init_population) * 100)
    summary += "\n" + "Total city quarantine days:".ljust(40, ' ') + "%.2f (%.2f %%)" % (total_CQ, total_CQ / float(days) / float(num_cities) * 100)
    summary += "\n" + "Avg. population in quarantine:".ljust(40, ' ') + "%.2f (%.2f %%)" % (total_quarantined / float(days) / float(num_cities), total_quarantined / float(days) / float(num_cities) / float(total_init_population) * 100)
    summary += "\n"
    return summary


def write_city_table_file(paths, file_name, headers, col_formats, col_vals, col_width=10):
    """Creates a table of values and writes it to the file(s) paths/file_name.
    :param paths: Directories the output file is written to
    :type paths: List of String
    :param file_name: File name
    :type file_name: String
    :param headers: List of header strings for each table column
    :type headers: List of String
    :param col_formats: List of formats per column (possible values are for instance "s", ".2f", "d")
    :type col_formats: String of String
    :param col_vals: Two-dimensional list of values per row per column
    :type col_vals: List of List
    :param col_width: Width of each column in characters
    :type col_width: int
    """

    num_cols = len(headers)
    num_rows = len(col_vals[0])

    for path in paths:
        outfile = open(os.path.join(path, file_name), 'w+')

        header_format = "\t".join(["%s"] * num_cols) + "\n"
        headers = tuple([header.rjust(col_width, ' ') for header in headers])
        outfile.write(header_format % headers)

        row_format = "\t".join(["%" + str(col_width) + col_format for col_format in col_formats]) + "\n"
        for row in range(0, num_rows):
            outfile.write(row_format % tuple([col[row] for col in col_vals]))

        outfile.close()


def make_plot(labels, col_city_vals, init_population, x):
    """Creates a plot with a subplot for each city.
    :param labels: Labels for each column
    :type labels: List of String
    :param col_city_vals: List of values per column per city
    :type col_city_vals: List of List of float
    :param init_population: Initial population (used for axes scale)
    :type: init_population: int
    :param x: Each column represents x days
    :type x: float (<= 1)
    :return: matplotlib.pyplot
    """

    # TODO: Add run to image

    num_cities = len(col_city_vals[0])
    num_cols = len(labels)

    gs = gridspec.GridSpec(num_cities, 1)

    # Adjust spaces between subplots
    plt.subplots_adjust(hspace=0.7, bottom=0.15)

    for city in range(0, num_cities):
        ax = plt.subplot(gs[city], title='City ' + str(city), xlabel='Days', ylabel='People')

        # Set axes limit
        ax.set_ylim([0, init_population * 1.5])

        # Adjust subplot title font size
        plt.setp(ax.title, fontsize=10)

        # Adjust axis labels and axis ticks font sizes
        ax.xaxis.label.set_size(10)
        ax.yaxis.label.set_size(10)
        plt.setp(ax.get_xticklabels(), rotation='horizontal', fontsize=10)
        plt.setp(ax.get_yticklabels(), rotation='horizontal', fontsize=10)

        # Set number of ticks
        ax.locator_params(axis='y', nbins=5)

        # Show the x axes label right of the axes
        ax.xaxis.set_label_coords(1.07, -0.1)

        # Shrink subplot width by 20% and move it 5 % to the right
        box = ax.get_position()
        ax.set_position([box.x0, box.y0, box.width * 0.8, box.height])
        for col in range(0, num_cols):
            ax.plot(x, col_city_vals[col][city], label=labels[col])

    ax.legend(loc='center left', bbox_to_anchor=(1.1, 0.5), fancybox=True, shadow=True, prop={'size': 10},
              labelspacing=0.2)

    return plt


def get_num_individuals_from_fractions(fraction):
    """This function uses a random generator to turn fractions of individuals into whole individuals.
    A fraction of x.y is converted into x + (y <= rand ? 1 : 0).
    :param fraction: Fraction of individuals
    :type fraction: float
    :return: Number of whole individuals
    """

    individuals = math.floor(fraction)
    probability = fraction - math.floor(fraction)
    if probability > 0 and random.random() <= probability:
        individuals += 1

    return individuals


def init_two_dim_list(len_dim1, len_dim2, init_val):
    """Initializes a two dimensional list with an initial value.
    :param len_dim1: Length of the first dimension
    :type len_dim1: int
    :param len_dim2: Length of the second dimension
    :type len_dim2: int
    :param init_val: Initial value
    """
    return [[init_val for _ in range(len_dim2)] for _ in range(len_dim1)]


sars()